# -*- coding: utf-8 -*-
import pandas as pd

from navec import Navec

path = 'navec_hudlit_v1_12B_500K_300d_100q.tar'
navec = Navec.load(path) #Загрузка коллекции эмбеддингов

import numpy as np
#Угол между векторами
def cosine_similarity(a, b):
    a = np.array(a)
    b = np.array(b)
    
    return np.dot(a, b) / np.linalg.norm(a) / np.linalg.norm(b)

import wikipedia
import random
import tqdm
import nltk
import re
import string
import threading
import multiprocessing
import os
import shutil
from nltk.corpus import stopwords

import numpy as np
import torch
from torch import nn

nltk.download('stopwords')
nltk.download('punkt')

wikipedia.set_lang("ru")
#Просмотр текстов из Википедии
def wikiTexts(topic, linksNum:list = [0], first:bool=True) -> (list, list) :
    links = []
    try:
        page = wikipedia.page(topic)
    except:
        return [], []
    
    try:
      links = page.links
    except:
      return result, [topic]
    
    result = [page.content]
    if len(linksNum) == 0:
        return result, [topic]
    
    random.shuffle(links)

    iterate = links[:linksNum[0]]
    resultLinks = [*iterate]
    if first:
      iterate = tqdm.tqdm_notebook(iterate)
    
    for link in iterate:
        contents = wikiTexts(link, linksNum[1:], False)
        result += contents[0]
        resultLinks += contents[1]
    return result, resultLinks

def wikiThread(returnDict, topic, linksNum:list = [0]):
  returnDict[topic] = wikiTexts(topic, linksNum)

depthSearch = [5, 2] # 5 первых ссылок, по 2 вторичные ссылки из каждой первой ссылки
jobs = []

manager = multiprocessing.Manager()
returnDict = manager.dict()

tasks = [
         (returnDict, "животные", depthSearch),
         (returnDict, "музыка", depthSearch),
         (returnDict, "спорт", depthSearch),
         (returnDict, "литература", depthSearch),
]
for i in range(len(tasks)):
  out_list = list()
  process = multiprocessing.Process(target=wikiThread, 
                                    args=tasks[i])
  jobs.append(process)

for j in jobs:
  j.start()

for j in jobs:
  j.join()
animalTexts, animalLinks = returnDict["животные"]
musicTexts, musicLinks = returnDict["музыка"]
sportTexts, sportLinks = returnDict["спорт"]
literatureTexts, literatureLinks = returnDict["литература"]

allTexts = (animalTexts, musicTexts, sportTexts, literatureTexts)
allLinks = (animalLinks, musicLinks, sportLinks, literatureLinks)
categories = ("животные", "музыка", "спорт", "литература")

for category, texts, links in zip(categories, allTexts, allLinks):
  print(category, ":", sum([len(text.split(' ')) for text in texts]), "слов")
  print(*links, sep=", ")

#Поиск слов определённой категории без эмбеддингов
#-------------------------------------------------
allTexts = [
            [[token.strip() for token in nltk.word_tokenize(text.lower())] for text in text_category] for text_category in allTexts
]

token_sets = {
    'животные': set(sum(allTexts[0], [])),
    'музыка': set(sum(allTexts[0], [])),
    'спорт': set(sum(allTexts[0], [])),
    'литература': set(sum(allTexts[0], [])),
}

from tqdm import tqdm
non_exclusive_words = set()
for word in tqdm(set.union(*token_sets.values())):
  count = 0
  for word_set in token_sets.values():
    if word in word_set:
      count+=1

  if count > 1:
    non_exclusive_words.add(word)

exclusive_word_sets = {name: word_set - non_exclusive_words for name, word_set in token_sets.items()}

stopWordsRu = set(stopwords.words("russian")) | set(["м", "см", "дм", "мм", "км", "руб", "коп", "ч", "мин", "с", "г", "кг", "т", "ц", "л", "мл"]) #Стоп-слова и единицы измерения
punctuation = set(string.punctuation + "-") #Пунктуация

def token_to_embed_id(word):
  #Если является стоп-словом или пунктуацией
  if word in stopWordsRu or word in punctuation or word.isdigit():
    return None
  
  #Если у слова есть эмбеддинг
  if word in navec.vocab:
    return navec.vocab[word]

  #Если слово встречается в текстах в Википедии
  for category_name, exclusive_word_set in exclusive_word_sets.items():
    if word in exclusive_word_set:
      return navec.vocab[category_name]

  return navec.vocab['<unk>']

#Идентификация текста как набора токенов
def text_to_ids(text: str):
  tokens = []
  embed_ids = []
  for word in nltk.word_tokenize(text):
    word = word.strip().lower()
    embed_id = token_to_embed_id(word)

    if embed_id is not None:
      tokens.append(word)
      embed_ids.append(embed_id)
  return tokens, embed_ids

from scipy.special import softmax

classes = [
    'животные',
    'спорт',
    'музыка',
    'литература'
]

classes_emb = [
               navec[class_name] for class_name in classes
]

math_emb = navec['математика']

#Получение вероятностей каждого класса
def get_sentence_predictions(sentence: str):
  tokens, embed_ids = text_to_ids(sentence)
  embeddings = [navec.pq[embed_id] for embed_id in embed_ids]

  results = [0, 0, 0, 0]
  for emb in embeddings:
    probs = [cosine_similarity(class_emb, emb) for class_emb in classes_emb]
    class_pred = np.argmax(probs)
    pred = probs[class_pred]
    if pred < 0 or cosine_similarity(math_emb, emb) > pred: continue #если ни одна категория не подходит или слово является математическим термином, пропустить
    
    probs[class_pred] = 0 #не учитывать определённую категорию
    if pred - max(probs) > 0.1: results[class_pred]+=pred - max(probs) #прибавить точность, если она выше 0.1
  return results

def get_significancy(sentence: str, class_: int):
  tokens, embed_ids = text_to_ids(sentence)
  embeddings = np.array([navec.pq[embed_id] for embed_id in embed_ids])

  results = []
  for token, emb in zip(tokens, embeddings):
    probs = [cosine_similarity(class_emb, emb) for class_emb in classes_emb]
    class_pred = np.argmax(probs)
    pred = probs[class_pred]
    if pred < 0 or cosine_similarity(math_emb, emb) > pred: continue #если ни одна категория не подходит или слово является математическим термином, пропустить

    probs[class_pred] = 0 #не учитывать определённую категорию
    if pred - max(probs) >= 0.1 and class_pred == class_: #если точность больше 0.1 и класс совпадает,
      results.append((pred - max(probs), token))
  
  return sorted(results, reverse=True)


import aiogram
import time
import logging
import config
from aiogram import types
import user_db
from aiogram.dispatcher.filters.state import State, StatesGroup
from aiogram.contrib.fsm_storage.memory import MemoryStorage
from aiogram.dispatcher import FSMContext

# Creating logger
logging.basicConfig(filename="WorkScrapper.log", level=logging.DEBUG)
logger = logging.getLogger("WorkScrapper")
# Initialization Bot
try:
    storage = MemoryStorage()
    bot = aiogram.Bot(config.API_TOKEN)
    dp = aiogram.Dispatcher(bot, storage = storage)
    logger.info(f"[{time.strftime('%D %H:%M:%S')}] Bot was initializated successfuly")
except Exception as e:
    print(e)
    logger.error(f"[{time.strftime('%D %H:%M:%S')}] Bot was not initializated due to error: {e}")

# Creating basic users database

users_db = user_db.Users_db("users.db")


# Listenning


class SearchOpts(StatesGroup):
    task = State()

# Welcome func
@dp.message_handler(commands=['start'])
async def welcome(message: types.Message):
    await message.answer(f"Здравствуйте, {message.from_user.username}! Для того, чтобы подобрать категорию задачи, введите /get_category")

# Help func
@dp.message_handler(commands=['help'])
async def help_func(message: types.Message):
    await message.answer("Вы можете воспользоваться этими командами:\n/start - Начать работу с ботом\n/delete_me - Удалить свой аккаунт\n/get_category - Определить категорию задачи")

# Change language func
@dp.message_handler(commands=['delete_me'])
async def delete_user(message: types.Message):
    req = users_db.delete_user(message.from_user.id)
    if req == "Ok":
        await message.answer(f"Вы удалили свой аккаунт ✅ Чтобы возобновить работу, используйте /start")
    else:
        logger.error(req)
        print(req)

# Search func
@dp.message_handler(commands=['get_category'])
async def enter_task(message: types.Message):
    await SearchOpts.task.set()
    await message.answer("Введите задачу: ")

@dp.message_handler(state=SearchOpts.task)
async def get_category(message: types.Message, state: FSMContext):
    preds = get_sentence_predictions(message.text)
    await message.answer("Предсказанная категория: " + classes[1 if preds[1] >= 0.2 else np.argmax(preds)])
    await state.finish()

if __name__ == '__main__':
	aiogram.executor.start_polling(dp, skip_updates=True)
